from fastapi import FastAPI, Request
from pydantic import BaseModel
import uvicorn
from loguru import logger
import time
import requests  # 用HTTP API替代Nacos SDK，稳定无兼容问题
from typing import Optional
from fastapi.responses import JSONResponse  # 新增导入
import traceback
# -------------------------- 服务初始化 --------------------------
app = FastAPI(title="RAG Distributed Agent Service")

# Nacos 配置（直接写死，简单粗暴不报错）
NACOS_SERVER = "http://192.168.150.101:8848"  # Nacos地址
SERVICE_NAME = "rag_agent_service"  # 服务名（和流程编排一致）
AGENT_IP = "192.168.150.1"  # 你的Windows机器IPv4（ipconfig查）
AGENT_PORT = 8007  # Agent服务端口
NACOS_NAMESPACE = ""  # public命名空间传空字符串


# -------------------------- 核心：用Nacos HTTP API注册服务（替代SDK） --------------------------
def register_to_nacos() -> bool:
    """调用Nacos官方HTTP API注册服务，无需SDK，100%兼容"""
    url = f"{NACOS_SERVER}/nacos/v1/ns/instance"
    params = {
        "serviceName": SERVICE_NAME,
        "ip": AGENT_IP,
        "port": AGENT_PORT,
        "weight": 10.0,
        "metadata": '{"version":"v1.0","type":"rag"}',  # JSON字符串
        "ephemeral": "true",  # 临时实例（服务下线自动剔除）
        "clusterName": "DEFAULT",
        "namespaceId": NACOS_NAMESPACE
    }
    try:
        # 发送HTTP请求注册
        response = requests.post(url, params=params, timeout=5)
        if response.status_code == 200 and response.text.strip() == "ok":
            logger.info(f" rag服务注册成功 code==200")
            return True
        else:
            logger.error(f"❌ 服务注册失败：Nacos返回 {response.status_code} - {response.text}")
            return False
    except Exception as e:
        logger.error(f"❌ 服务注册异常：{str(e)}（检查Nacos地址是否可达、Windows防火墙是否开放）")
        return False


# # 启动时执行注册（失败直接退出）
# if not register_to_nacos():
#     exit(1)


# -------------------------- 请求/响应模型 --------------------------
class AgentRequest(BaseModel):
    flow_id: str
    query: str
    trace_id: str


class AgentResponse(BaseModel):
    code: int
    data: str
    message: str
    trace_id: str


# -------------------------- 中间件（链路追踪） --------------------------
@app.middleware("http")
async def add_trace_middleware(request: Request, call_next):
    # 获取或生成trace_id
    trace_id = request.headers.get("X-Trace-ID", f"trace_agent_{str(time.time_ns())[:16]}")

    # 关键修复：用bind()添加上下文，而非configure()修改全局配置
    logger.bind(trace_id=trace_id).info(f"[Agent请求] path={request.url.path}，trace_id={trace_id}")

    # 处理请求
    response = await call_next(request)
    response.headers["X-Trace-ID"] = trace_id
    return response


# -------------------------- RAG核心接口 --------------------------
@app.post(
    "/agent/rag/invoke", 
    response_model=AgentResponse,
    response_class=JSONResponse#,
    # media_type="application/json"  # 明确指定媒体类型
)
async def invoke_rag_agent(req: AgentRequest):
    try:
        # 修复结果格式，避免使用多行字符串导致的序列化问题
        rag_result = (
            f"RAG Agent 处理结果（trace_id: {req.trace_id}）：\n"
            f"- 查询内容：{req.query}\n"
            "- 匹配知识库：分布式缓存技术手册（v2.3）\n"
            "- 核心结论：分布式缓存常用Redis Cluster，支持分片存储，通过主从复制实现高可用，需注意缓存穿透（用布隆过滤器）、缓存雪崩（设置随机过期时间）问题。"
        )
        logger.info(f"[RAG处理成功] flow_id={req.flow_id}")
        return {
            "code": 200,
            "data": rag_result.strip(),
            "message": "success",
            "trace_id": req.trace_id
        }
    except Exception as e:
        err_msg = f"RAG处理失败：{str(e)}"
        logger.error(f"[RAG处理异常] flow_id={req.flow_id}，err={err_msg}，堆栈：{traceback.format_exc()}")  # 新增堆栈信息
        return {
            "code": 500,
            "data": "",
            "message": err_msg,
            "trace_id": req.trace_id
        }


# -------------------------- 健康检查接口（Nacos心跳用） --------------------------
@app.get("/health")
async def health_check():
    return {"status": "healthy", "service": SERVICE_NAME, "port": AGENT_PORT}


# -------------------------- 启动入口 --------------------------
# ✅ 正确：只在直接运行时注册
if __name__ == "__main__":
    if not register_to_nacos():
        print("❌ RAG 服务注册失败")
        exit(1)
    uvicorn.run(app, host="0.0.0.0", port=AGENT_PORT, log_config=None)